#!/usr/bin/env python3

# PPU Assembler

# This is really just intended for testing; you need to generate PPU code
# dynamically if you want animated scenes. Example program:
#
# .define test_img_loc 0x20010000
#
#     clip 0, 319
#
# loop:
#     fill 31, 0, 16
#
#     blit 50, 100, size=SIZE_256, img=test_img_loc, fmt=FORMAT_ARGB1555
#
#     ablit 0, 0, size=SIZE_256, img=test_img_loc, fmt=FORMAT_ARGB1555,\
#         a00=0x000, a01=0x100,\
#         a10=0x100, a11=0x000
#
#     sync
#     push loop
#     popj

import re
import struct
import sys

# Each instruction consists of one or more words. For each instruction we
# provide here a list of words, each consisting of a list of (name/const,
# lshift, mask lsb, mask msb) for the operands encoded in that instruction
# word. The mask is inclusive at both ends.
instruction_encoding = {
	"sync": (
		(
			(0x0,        28, 28, 31),
		),
	),
	"clip": (
		(
			(0x1,        28, 28, 31),
			("start",    0,  0,  9 ),
			("end",      10, 10, 19),
		),
	),
	"fill": (
		(
			(0x2,        28, 28, 31),
			("r",        10, 10, 14),
			("g",        5,  5,  9 ),
			("b",        0,  0,  4 ),
		),
	),
	"blit": (
		(
			(0x4,        28, 28, 31),
			("x",        0,  0,  9 ),
			("y",        10, 10, 19),
			("poff",     22, 22, 24),
			("size",     25, 25, 27),
		),
		(
			("img",      0,  2,  31),
			("fmt",      0,  0,  1 ),
		),
	),
	"tile": (
		(
			(0x5,        28, 28, 31),
			("x",        0,  0,  9 ),
			("y",        10, 10, 19),
			("poff",     22, 22, 24),
			("size",     25, 25, 26),
		),
		(
			("tilemap",  0,  2,  31),
			("pfsize",   0,  0,  1 ),
		),
		(
			("tileset",  0,  2,  31),
			("fmt",      0,  0,  1 ),
		),
	),
	"ablit": (
		(
			(0x6,        28, 28, 31),
			("x",        0,  0,  9 ),
			("y",        10, 10, 19),
			("halfsize", 21, 21, 21),
			("poff",     22, 22, 24),
			("size",     25, 25, 27),
		),
		(
			("b0",       0,  0,  15),
			("b1",       16, 16, 31),
		),
		(
			("a00",      0,  0,  15),
			("a01",      16, 16, 31),
		),
		(
			("a10",      0,  0,  15),
			("a11",      16, 16, 31),
		),
		(
			("img",      0,  2,  31),
			("fmt",      0,  0,  1 ),
		),
	),
	"atile": (
		(
			(0x7,        28, 28, 31),
			("x",        0,  0,  9 ),
			("y",        10, 10, 19),
			("poff",     22, 22, 24),
			("size",     25, 25, 26),
		),
		(
			("b0",       0,  0,  15),
			("b1",       16, 16, 31),
		),
		(
			("a00",      0,  0,  15),
			("a01",      16, 16, 31),
		),
		(
			("a10",      0,  0,  15),
			("a11",      16, 16, 31),
		),
		(
			("tilemap",  0,  2,  31),
			("pfsize",   0,  0,  1 ),
		),
		(
			("tileset",  0,  2,  31),
			("fmt",      0,  0,  1 ),
		),
	),
	"push": (
		(
			(0xe,        28, 28, 31),
		),
		(
			("target",   0,  0,  31),
		),
	),
	"popj": (
		(
			(0xf,        28, 28, 31),
			("cond",     24, 24, 27),
			("compval",  0,  0,  9 ),
		),
	),
}


# Our instructions have a lot of arguments, so we allow them
# to be passed positionally e.g.        clip 0 319
# or keyword-wise with defaults e.g.    fill r=31 b=16
# This table specifies the arguments, their positional order,
# and default if any.
instruction_arg_keywords = {
	"sync": (),
	"clip": (
		("start", None),
		("end", None)),
	"fill": (
		("r", 0),
		("g", 0),
		("b", 0)),
	"blit": (
		("x", 0),
		("y", 0),
		("size", 0),
		("poff", 0),
		("fmt", 0),
		("img", None)),
	"ablit": (
		("x", 0),
		("y", 0),
		("size", 0),
		("halfsize", 0),
		("poff", 0),
		("fmt", 0),
		("img", None),
		("a00", 1 << 8),
		("a01", 0 << 8),
		("a10", 0 << 8),
		("a11", 1 << 8),
		("b0", 0),
		("b1", 0)),
	"tile": (
		("x", 0),
		("y", 0),
		("size", 0),
		("poff", 0),
		("pfsize", 0),
		("tilemap", None),
		("fmt", 0),
		("tileset", None)),
	"atile": (
		("x", 0),
		("y", 0),
		("size", 0),
		("poff", 0),
		("pfsize", 0),
		("tilemap", None),
		("fmt", 0),
		("tileset", None),
		("a00", 1 << 8),
		("a01", 0 << 8),
		("a10", 0 << 8),
		("a11", 1 << 8),
		("b0", 0),
		("b1", 0)),
	"push": (
		("target", None),),
	"popj": (
		("cond", 0),
		("compval", 0)),
}

default_defs = {
	"SIZE_8":          0,
	"SIZE_16":         1,
	"SIZE_32":         2,
	"SIZE_64":         3,
	"SIZE_128":        4,
	"SIZE_256":        5,
	"SIZE_512":        6,
	"SIZE_1024":       7,
	"FORMAT_ARGB1555": 0,
	"FORMAT_PAL8":     1,
	"FORMAT_PAL4":     2,
	"FORMAT_PAL1":     3,
	"COND_YLT":        1,
	"COND_YGE":        2,
	"ylt":             1,
	"yge":             2,
}

def getint(x):
	if type(x) is int:
		return x # "can't convert non-string with explicit base"
	else:
		return int(x, 0)

def resolve_defs(val, defs):
	stack = []
	while val in defs and val not in stack:
		stack.append(val)
		val = defs[val]
	return val

def resolve_instr_args(linenum, instr, args, defs):
	# Tuple of (tuples of (name, default value)):
	argspec = instruction_arg_keywords[instr]
	if len(args) > len(argspec):
		sys.exit(f"Line {linenum + 1}: too many operands for instruction {instr}")
	resolved_args = dict(argspec)
	for i, arg in enumerate(args):
		if "=" in arg:
			parts = arg.split("=")
			if len(parts) != 2:
				sys.exit(f"Line {linenum + 1}: ill-formed keyword operand '{arg}'")
			k, v = parts
			if k not in resolved_args:
				sys.exit(f"Line {linenum + 1}: unknown operand {k} for instruction {instr}")
			resolved_args[k] = resolve_defs(v, defs)
		else:
			resolved_args[argspec[i][0]] = resolve_defs(arg, defs)
	for k, v in resolved_args.items():
		if v is None:
			sys.exit(f"Line {linenum + 1}: non-optional operand '{k}' of instruction '{instr}' was left unspecified")
	return resolved_args

if __name__ == "__main__":
	if len(sys.argv) < 2 or len(sys.argv) > 3:
		sys.exit("Usage: ppuasm <input.S> [output.bin]")

	carried_over_from_prev = None
	defs = default_defs.copy()
	labels = {}
	label_fixups = []
	program = []

	for linenum, line in enumerate(open(sys.argv[1])):
		line = line.strip()
		if carried_over_from_prev is not None:
			line = " ".join((carried_over_from_prev, line))
			carried_over_from_prev = None
		line = line.split("#")[0].strip()
		if len(line) == 0:
			continue
		if line.endswith("\\"):
			carried_over_from_prev = line[:-1]
			continue
		tokens = re.split(r"\s*[,\s]\s*", line)

		if tokens[0] == ".define":
			if len(tokens) != 3:
				sys.exit(f"Line {linenum + 1}: .define requires 2 arguments")
			defs[tokens[1]] = tokens[2]
		elif tokens[0].endswith(":"):
			if len(tokens) != 1:
				sys.exit(f"Line {linenum + 1}: garbage after label {tokens[0]}")
			labels[tokens[0][:-1]] = len(program) * 4
		elif tokens[0].lower() in instruction_arg_keywords:
			instr = tokens[0].lower()
			args = resolve_instr_args(linenum, instr, tokens[1:], defs)
			if instr == "push":
				label_fixups.append((linenum, len(program) + 1, args["target"]))
				args["target"] = 0
			for word_fmt in instruction_encoding[instr]:
				word = 0
				for fieldspec in word_fmt:
					if type(fieldspec[0]) is int:
						val = fieldspec[0]
					else:
						try:
							argval = args[fieldspec[0]]
							val = getint(argval)
						except ValueError:
							sys.exit(f"Line {linenum + 1}: '{argval}' is not a valid integer")
					word |= (val << fieldspec[1]) & ((1 << fieldspec[3] + 1) - (1 << fieldspec[2]))
				program.append(word)
		else:
			sys.exit(f"Line {linenum + 1}: invalid instruction '{tokens[0]}'")


	for linenum, fixloc, labelname in label_fixups:
		if labelname not in labels:
			sys.exit(f"Line {linenum + 1}: label '{labelname}' is not defined in program")
		program[fixloc] = labels[labelname]

	if len(sys.argv) > 2:
		with open(sys.argv[2], "wb") as ofile:
			ofile.write(struct.pack("<" + "L" * len(program), *program))
	else:
		for i, w in enumerate(program):
			print(f"{4*i:04x}: {w:08x}")
